#include "kernel_operator.h"
#include <vector>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelCross
{
public:
  __aicore__ inline KernelCross() {}
  __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, uint32_t iValue, uint32_t jValue, uint32_t x1iValue, uint32_t x1jValue, uint32_t x2iValue, uint32_t x2jValue, 
                              uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength,
                              uint32_t typeKey, uint32_t totalLength, uint32_t dim, uint32_t dimNum, uint32_t * x1shape, uint32_t * x2shape, uint32_t * yshape)
  {
    this->typeKey = typeKey;
    this->dim = dim;
    this->dimNum = dimNum;
    this->x1shape = x1shape;
    this->x2shape = x2shape;
    this->yshape = yshape;
    this->iValue = iValue;
    this->x1iValue = x1iValue;
    this->x2iValue = x2iValue;
    this->jValue = jValue;
    this->x1jValue = x1jValue;
    this->x2jValue = x2jValue;
    this->totalLength = totalLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;
    this->x1_addr = x1;
    this->x2_addr = x2;

    yGm.SetGlobalBuffer((__gm__ DTYPE_X1 *)y,
                        this->totalLength);
    pipe.InitBuffer(inQueue, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
  }
  __aicore__ inline void Process()
  {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++)
    {
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
  }

private:
  __aicore__ inline void CopyIn(int32_t progress)
  {
    LocalTensor<DTYPE_X1> inLocal = inQueue.AllocTensor<DTYPE_X1>();
    inQueue.EnQue(inLocal);
  }
  __aicore__ inline void Compute(int32_t progress)
  {
    LocalTensor<DTYPE_X1> inLocal = inQueue.DeQue<DTYPE_X1>();
    LocalTensor<DTYPE_X1> yLocal = outQueueY.AllocTensor<DTYPE_X1>();
    
    uint32_t index = progress * this->tileLength;
    if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
        (progress == (this->tileNum * BUFFER_NUM - 1)))
    {
      index = (progress - 2) * (this->tileLength) + this->lasttileLength;
    }
    for (uint32_t p = 0; p < this->tileLength; p++, index++)
    {
      if (index < 0)
      {
        continue;
      }
      uint32_t yi = GetIIndex(index);
      uint32_t yj = GetJIndex(index);
      uint32_t k = GetKIndex(index);
      uint32_t x1i = GetX1IIndex(yi);
      uint32_t x2i = GetX2IIndex(yi);
      uint32_t x1j = GetX1JIndex(yj);
      uint32_t x2j = GetX2JIndex(yj);
      if (k == 0)
      {
        uint32_t x1Index = GetX1Index(x1i, 1, x1j);
        uint32_t x2Index = GetX2Index(x2i, 2, x2j);
        float ans = ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]);
        x1Index = GetX1Index(x1i, 2, x1j);
        x2Index = GetX2Index(x2i, 1, x2j);
        ans = ((float)ans - ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]));
        yLocal.SetValue(p, (DTYPE_X1)ans);
      }
      else if (k == 1)
      {
        uint32_t x1Index = GetX1Index(x1i, 2, x1j);
        uint32_t x2Index = GetX2Index(x2i, 0, x2j);
        float ans = ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]);
        x1Index = GetX1Index(x1i, 0, x1j);
        x2Index = GetX2Index(x2i, 2, x2j);
        ans = ((float)ans - ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]));
        yLocal.SetValue(p, (DTYPE_X1)ans);
      }
      else
      {
        uint32_t x1Index = GetX1Index(x1i, 0, x1j);
        uint32_t x2Index = GetX2Index(x2i, 1, x2j);
        float ans = ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]);
        x1Index = GetX1Index(x1i, 1, x1j);
        x2Index = GetX2Index(x2i, 0, x2j);
        ans = ((float)ans - ((float)((__gm__ DTYPE_X1 *)x1_addr)[x1Index] * (float)((__gm__ DTYPE_X1 *)x2_addr)[x2Index]));
        yLocal.SetValue(p, (DTYPE_X1)ans);
      }
    }
    outQueueY.EnQue(yLocal);
    inQueue.FreeTensor(inLocal);
  }

  __aicore__ inline void CopyOut(int32_t progress)
  {
    LocalTensor<DTYPE_X1> yLocal = outQueueY.DeQue<DTYPE_X1>();
    if (BUFFER_NUM == 1)
    {
      if (progress == this->tileNum - 1)
      {
        if (progress == 0)
        {
          DataCopy(yGm[0], yLocal, this->tileLength);
        }
        else
        {
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              yLocal, this->tileLength);
        }
      }
      else
      {
        DataCopy(yGm[progress * this->tileLength], yLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2)
    {
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1)))
      {
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            yLocal, (this->tileLength));
      }
      else
      {
        DataCopy(yGm[progress * (this->tileLength)], yLocal,
                 (this->tileLength));
      }
    }
    outQueueY.FreeTensor(yLocal);
  }
  __aicore__ inline uint32_t GetIndex(uint32_t i, uint32_t k, uint32_t j)
  {
    return i * (3 * this->jValue) + k * jValue + j;
  }
  __aicore__ inline uint32_t GetX1Index(uint32_t i, uint32_t k, uint32_t j)
  {
    return i * (3 * this->x1jValue) + k * this->x1jValue + j;
  }
  __aicore__ inline uint32_t GetX2Index(uint32_t i, uint32_t k, uint32_t j)
  {
    return i * (3 * this->x2jValue) + k * this->x2jValue + j;
  }
  __aicore__ inline uint32_t GetIIndex(uint32_t index)
  {
    return index / (3 * this->jValue);
  }
  __aicore__ inline uint32_t GetJIndex(uint32_t index)
  {
    return index % this->jValue;
  }
  __aicore__ inline uint32_t GetKIndex(uint32_t index)
  {
    return (index / this->jValue) % 3;
  }

  __aicore__ inline uint32_t GetX1IIndex(uint32_t yi)
  {
    uint32_t num = 1;
    uint32_t ans = 0;
    for (int i = dim - 1; i >= 0; i--)
    {
      if (yshape[i] == x1shape[i])
      {
        ans += (yi % yshape[i]) * num;
      }
      num *= x1shape[i];
      yi /= yshape[i];
      if (yi == 0)
        break;
    }
    return ans;
  }
  __aicore__ inline uint32_t GetX1JIndex(uint32_t yj)
  {
    uint32_t num = 1;
    uint32_t ans = 0;
    for (int i = dimNum - 1; i > dim; i--)
    {
      if (yshape[i] == x1shape[i])
      {
        ans += (yj % yshape[i]) * num;
      }
      num *= x1shape[i];
      yj /= yshape[i];
      if (yj == 0)
        break;
    }
    return ans;
  }

  __aicore__ inline uint32_t GetX2IIndex(uint32_t yi)
  {
    uint32_t num = 1;
    uint32_t ans = 0;
    for (int i = dim - 1; i >= 0; i--)
    {
      if (yshape[i] == x2shape[i])
      {
        ans += (yi % yshape[i]) * num;
      }
      num *= x2shape[i];
      yi /= yshape[i];
      if (yi == 0)
        break;
    }
    return ans;
  }
  __aicore__ inline uint32_t GetX2JIndex(uint32_t yj)
  {
    uint32_t num = 1;
    uint32_t ans = 0;
    for (int i = dimNum - 1; i >= 0; i--)
    {
      if (yshape[i] == x2shape[i])
      {
        ans += (yj % yshape[i]) * num;
      }
      num *= x2shape[i];
      yj /= yshape[i];
      if (yj == 0)
        break;
    }
    return ans;
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueue;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
  GlobalTensor<DTYPE_X1> yGm;
  GM_ADDR x1_addr;
  GM_ADDR x2_addr;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t totalLength;
  uint32_t lasttileLength;
  uint32_t typeKey;
  uint32_t iValue;
  uint32_t jValue;
  uint32_t x1iValue;
  uint32_t x1jValue;
  uint32_t x2iValue;
  uint32_t x2jValue;
  uint32_t dim;
  uint32_t dimNum;
  uint32_t * yshape;
  uint32_t * x1shape;
  uint32_t * x2shape;
};

extern "C" __global__ __aicore__ void cross(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
  GET_TILING_DATA(tiling_data, tiling);
  KernelCross op;

  op.Init(x1, x2, y, tiling_data.iValue, tiling_data.jValue, tiling_data.x1iValue, tiling_data.x1jValue, tiling_data.x2iValue, tiling_data.x2jValue, 
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, tiling_data.typeKey, tiling_data.totalLength, 
          tiling_data.dim, tiling_data.dimNum, tiling_data.x1shape, tiling_data.x2shape, tiling_data.yshape);
  op.Process();
}

#ifndef __CCE_KT_TEST__
void cross_do(uint32_t blockDim, void *l2ctrl, void *stream,
              uint8_t *x1, uint8_t *x2, uint8_t *y,
              uint8_t *workspace, uint8_t *tiling)
{
  cross<<<blockDim, l2ctrl, stream>>>(x1, x2, y, workspace, tiling);
}
#endif